import random
import json

# 设置整数参数范围
displacement_range = (-200, 200)  # mm，整数
rotation_range = (-30, 30)        # degrees，整数
components = ['x', 'y', 'z', 'rot_x', 'rot_y', 'rot_z']

# 自然语言模板
displacement_templates = {
    'x': "沿 x 轴运动 {v}毫米",
    'y': "沿 y 轴运动 {v}毫米",
    'z': "沿 z 轴运动 {v}毫米"
}
rotation_templates = {
    'rot_x': "绕 x 轴转动 {v}度",
    'rot_y': "绕 y 轴转动 {v}度",
    'rot_z': "绕 z 轴转动 {v}度"
}

def generate_sample():
    sample = {'x': 0, 'y': 0, 'z': 0, 'rot_x': 0, 'rot_y': 0, 'rot_z': 0}
    text_segments = []

    # 随机选择运动分量的个数（1~6）
    num_components = random.randint(1, 6)
    chosen_components = random.sample(components, num_components)

    # 随机顺序排列
    random.shuffle(chosen_components)

    for comp in chosen_components:
        if comp.startswith('rot'):
            value = random.randint(*rotation_range)  # 整数角度
            sample[comp] = value
            text_segments.append(rotation_templates[comp].format(v=value))
        else:
            value = random.randint(*displacement_range)  # 整数位移
            sample[comp] = value
            text_segments.append(displacement_templates[comp].format(v=value))

    question = "，".join(text_segments) + "。"
    response = json.dumps(sample, ensure_ascii=False)

    return {
        "Question": question,
        "Complex_CoT": "",
        "Response": response
    }

# 生成多个样本
def generate_dataset(n):
    return [generate_sample() for _ in range(n)]

# 保存为 jsonl 文件
def save_to_jsonl(dataset, filename="stewart_dataset.json"):
    with open(filename, 'w', encoding='utf-8') as f:
        for item in dataset:
            f.write(json.dumps(item, ensure_ascii=False) + '\n')

# 示例：生成5000条数据并保存
if __name__ == "__main__":
    dataset = generate_dataset(5000)
    save_to_jsonl(dataset)
    print("✅ 数据集已保存为 stewart_dataset.json（共 5000 条）")
